#pragma once
#include <opencv2/opencv.hpp>
#include <chrono>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <vector>
#include <map>
#include <cstring>
#include <fstream>
#include <numeric>
#include <algorithm>

#include "pdmodel_error_code.h"
#include "pdmodel_struct.h"
#include "model_manager.h"
#include "common/base_struct.h"

class LogInfo;

namespace ai {

class ModelConfig;
class PDClsPredictor;
class PDOCRRecognizer;
class PDOCRDetector;
class PDDetPredictor;
class PDSegPredictor;
class PDOCRPostProcessor;

class PDModelManager: public ai::ModelManager
{
public:
    PDModelManager(LogInfo *log_info);
    ~PDModelManager();

    int init_model(std::vector<InitModelData> &init_model_datas);

    int rec_img(std::string algo_type,
                    cv::Mat& img, 
                    std::vector<BaseInfo*> &data_info);

    // 线程执行
    int stop_run_thread();
    int start_run_thread(int cam_size);
    int set_data(int idx, cv::Mat& img);
    int set_data(int idx, stream::ImageBlob& img) {return -1;}
    int get_data(int idx, std::vector<BaseInfo*>& data);

    int free_model();
    int text_detect(cv::Mat& img, 
                    std::vector <BaseInfo*> &text_lines, 
                    int mdl_idx,
                    bool sort_by_default = true);

    // img 矫正后的图片
    int text_rec_lite(cv::Mat& img, 
                std::vector<BaseInfo*> &rec_ifo, 
                int mdl_idx );
    // img 原始图片 直接ocr
    int text_recognise(cv::Mat& img, 
                        TextDetInfo &det_info,
                        TextRecInfo &crnn_info, 
                        int model_idx,
                        float unclip_ratio = 0.);
    // img 原始图片 先det 再ocr
    int text_recognise(cv::Mat& img, 
                  std::vector<BaseInfo*> &rec_ifo, 
                  int mdl_idx );
                  
    int plate_rec(cv::Mat& img,
                std::vector<BaseInfo*>& cls_info, 
                int mdl_idx);
    int cls(cv::Mat& img,
                std::vector<BaseInfo*>& cls_info, 
                int mdl_idx);
    int cls(std::vector<cv::Mat> &imgs,
                std::vector<BaseInfo*>& cls_ifos, 
                int mdl_idx);

    int seg(cv::Mat& img, 
                std::vector<BaseInfo*> &seg_info,
                int mdl_idx);

    int det(cv::Mat& img, 
                std::vector<BaseInfo*> &det_info,
                int mdl_idx);

    int det(std::vector<cv::Mat> &imgs,
              std::vector<std::vector<DetInfo>>& det_infos, 
              int mdl_idx);
private:

    int load_model(const std::string &model_dir, 
                    ModelType mt,
                    ModelConfig* model_config,
                    paddle_infer::Config &config) ;

public:
    std::map<std::string, ModuleInfo> total_mdl;
    std::map<int, ModelConfig*> configs;
    int init_ok;

private:
    // friend class ClsPredictor;
    // friend class OCRRecognizer;
    // friend class OCRDetector;
    friend class PDDetPredictor;
    // friend class SegPredictor;

    std::map<int, std::unique_ptr<paddle::PaddlePredictor> > predictors;

    // PDClsPredictor * cls_pd_;
    // PDOCRDetector *text_det_;
    // PDOCRRecognizer *text_rec_;
    PDDetPredictor *det_pd_;
    // PDSegPredictor *seg_pd_;
    // PDOCRPostProcessor *ocr_pp;


    std::atomic_int stop_flag1;
    std::atomic_int stop_flag2;

};

}
